package com.well.kernel.spring.jdbc;

import org.apache.log4j.Logger;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Copyright &copy; Well All rights reserved.
 * Author:Well
 * Date:2017-10-25
 * Description:动态数据源
 */
public class DynamicDataSource extends AbstractRoutingDataSource
{
    private static Logger logger=Logger.getLogger(DynamicDataSource.class);

    private DataSource master; // 主库，只允许有一个
    private List<DataSource> slaves; // 从库，允许有多个
    private AtomicInteger slaveCount = new AtomicInteger();
    private int slaveSize = 0;

    private Map<Object, Object> dataSources = new HashMap<Object, Object>();

    private static final String DEFAULT = "master";
    private static final String SLAVE = "slave";

    private static final ThreadLocal<LinkedList<String>> datasourceHolder = new ThreadLocal<LinkedList<String>>() {

        @Override
        protected LinkedList<String> initialValue() {
            return new LinkedList<String>();
        }

    };

    /**
     * 初始化
     */
    @Override
    public void afterPropertiesSet() {
        if (null == master) {
            throw new IllegalArgumentException("Property 'master' is required");
        }
        dataSources.put(DEFAULT, master);
        if (null != slaves && slaves.size() > 0) {
            for (int i = 0; i < slaves.size(); i++) {
                dataSources.put(SLAVE + (i + 1), slaves.get(i));
            }
            slaveSize = slaves.size();
        }
        this.setDefaultTargetDataSource(master);
        this.setTargetDataSources(dataSources);
        super.afterPropertiesSet();
    }
    /**
     * 选择使用主库，并把选择放到当前ThreadLocal的栈顶
     */
    public static void useMaster() {
        logger.debug("use master data source ");
        LinkedList<String> m = datasourceHolder.get();
        m.offerFirst(DEFAULT);
    }

    /**
     * 选择使用从库，并把选择放到当前ThreadLocal的栈顶
     */
    public static void useSlave() {
        logger.debug("use slave data source ");
        LinkedList<String> m = datasourceHolder.get();
        m.offerFirst(SLAVE);
    }

    /**
     * 重置当前栈
     */
    public static void reset() {
        LinkedList<String> m = datasourceHolder.get();
        if (m.size() > 0) {
            logger.debug("reset data source " + m);
            m.poll();
        }else{
            logger.debug("reset master data source");
        }
    }

    /**
     * 如果是选择使用从库，且从库的数量大于1，则通过取模来控制从库的负载,
     * 计算结果返回AbstractRoutingDataSource
     */
    @Override
    protected Object determineCurrentLookupKey() {
        LinkedList<String> m = datasourceHolder.get();
        String key=m.peekFirst();
        if(null==key || DEFAULT.equals(key) || slaveSize==0){
            return DEFAULT;
        }
        if(slaveSize==1){
            return SLAVE+"1";
        }
        Integer index = slaveCount.incrementAndGet()%slaveSize;
        if(slaveCount.get()>10000){
            slaveCount.set(-1);
        }
        return SLAVE + (index + 1);

        /*String key = m.peekFirst() == null ? DEFAULT : m.peekFirst();
        logger.debug("current datasource :" + key);
        if(DEFAULT.equals(key)){
            return key;
        }
        if(slaveSize > 1){
            long c = slaveCount.incrementAndGet();
            c = c % slaveSize;
            return SLAVE + (c + 1);
        }
        return SLAVE + "1";*/

        /*if (null != key) {
            if (DEFAULT.equals(key)) {
                return key;
            } else if (SLAVE.equals(key)) {
                if (slaveSize > 1) {// Slave loadBalance
                    long c = slaveCount.incrementAndGet();
                    c = c % slaveSize;
                    return SLAVE + (c + 1);
                } else {
                    return SLAVE + "1";
                }
            }
            return null;
        } else {
            return null;
        }*/
    }

    public DataSource getMaster() {
        return master;
    }

    public List<DataSource> getSlaves() {
        return slaves;
    }

    public void setMaster(DataSource master) {
        this.master = master;
    }

    public void setSlaves(List<DataSource> slaves) {
        this.slaves = slaves;
    }
}
